import os
import shutil
import zipfile
import urllib.request
def download_cornell_movie_dialogs(data_dir="cornell_movie_dialogs"):
    """
    下载并解压Cornell Movie-Dialogs数据集（PyTorch标准流程）
    官方源：https://www.cs.cornell.edu/~cristian/data/cornell_movie_dialogs_corpus.zip
    """
    os.makedirs(data_dir, exist_ok=True)
    zip_path = os.path.join(data_dir, "cornell_movie_dialogs_corpus.zip")
    
    # 下载数据集
    if not os.path.exists(zip_path):
        print("正在下载Cornell Movie-Dialogs数据集...")
        url = "https://www.cs.cornell.edu/~cristian/data/cornell_movie_dialogs_corpus.zip"
        urllib.request.urlretrieve(url, zip_path)
        print("下载完成")
    
    # 解压
    extract_dir = os.path.join(data_dir, "extracted")
    if not os.path.exists(extract_dir):
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)
        # 移动文件到上层目录
        for item in os.listdir(os.path.join(extract_dir, "cornell movie-dialogs corpus")):
            shutil.move(
                os.path.join(extract_dir, "cornell movie-dialogs corpus", item),
                os.path.join(extract_dir, item)
            )
    
    return extract_dir


def load_cornell_conversations(data_dir):
    """加载对话数据（遵循PyTorch数据加载规范）"""
    # 加载所有对话行
    lines = {}
    lines_path = os.path.join(data_dir, "movie_lines.txt")
    with open(lines_path, encoding="iso-8859-1") as f:
        for line in f:
            parts = line.strip().split(" +++$+++ ")
            if len(parts) == 5:
                line_id, _, _, _, text = parts
                lines[line_id] = text.strip()
    
    # 加载对话序列
    conversations = []
    conv_path = os.path.join(data_dir, "movie_conversations.txt")
    with open(conv_path, encoding="iso-8859-1") as f:
        for line in f:
            parts = line.strip().split(" +++$+++ ")
            if len(parts) == 4:
                _, _, _, line_ids = parts
                line_ids = eval(line_ids)  # 转换为列表
                # 提取完整对话
                conv = [lines[line_id] for line_id in line_ids if line_id in lines]
                if len(conv) >= 2:  # 至少两轮对话
                    conversations.append(conv)
    
    print(f"成功加载 {len(conversations)} 个多轮对话")
    return conversations



if __name__ == "__main__":
    # data_dir = download_cornell_movie_dialogs()
    
    data_dir = r'C:\Users\COLORFUL\Desktop\AI_NLP\transform框架\生成式对话\dataset\cornell_movie_dialogs\extracted'
    
    conversations = load_cornell_conversations(data_dir)
    print(conversations)